from contextvars import ContextVar
import logging
import uuid

from fastapi import Request

# 用于存储当前请求的request_id
_request_id_var = ContextVar('request_id', default=None)

class RequestIdFilter(logging.Filter):
    def filter(self, record):
        record.request_id = _request_id_var.get() or "NO_REQUEST_ID"
        return True

async def request_id_middleware(request: Request, call_next):
    request_id = str(uuid.uuid4())
    token = _request_id_var.set(request_id)
    try:
        response = await call_next(request)
        if response:
            response.headers["X-Request-ID"] = request_id
        return response
    finally:
        _request_id_var.reset(token)